#!/bin/bash

set -e

root=$(cd $(dirname ${0})/..; pwd)
build=${root}/build
node_modules=${root}/node_modules
src=${root}/src
test=${root}/test
third_party=${root}/third_party
tools=${root}/tools

identifier=pytorch
virtualenv=${third_party}/virtualenv/${identifier}

bold() {
    echo "$(tty -s && tput bold)$1$(tty -s && tput sgr0)" 
}

git_sync() {
    mkdir -p "${third_party}"
    if [ -d "${third_party}/${1}" ]; then
        git -C "${third_party}/${1}" pull --quiet --prune
    else
        git -C "${third_party}" clone --quiet --recursive ${2} ${1}
    fi
    git -C "${third_party}/${1}" submodule sync --quiet
    git -C "${third_party}/${1}" submodule update --quiet --init --recursive
}

clean() {
    bold "pytorch clean"
    rm -rf ${virtualenv}
    rm -rf ${third_party}/${identifier}
}

sync() {
    bold "pytorch sync"
    git_sync pytorch https://github.com/pytorch/pytorch.git
    git_sync torchvision https://github.com/pytorch/vision.git
}

install() {
    bold "pytorch install"
    # case "$(uname)" in
    #     "Linux")
    #         [ -n "$(which cmake)" ] || sudo apt install -y cmake
    #         ;;
    #     "Darwin")
    #         brew list automake > /dev/null 2>&1 || brew install automake > /dev/null
    #         brew list cmake > /dev/null 2>&1 || brew install cmake > /dev/null
    #         brew list gflags > /dev/null 2>&1 || brew install gflags > /dev/null
    #         brew list glog > /dev/null 2>&1 || brew install glog > /dev/null
    #         brew list protobuf > /dev/null 2>&1 && brew uninstall protobuf > /dev/null
    #         ;;
    # esac
    [ -n "$(python3 -m pip list --format columns --disable-pip-version-check | grep -w virtualenv)" ] || python3 -m pip install --force-reinstall --user --quiet virtualenv
    [ -d "${virtualenv}" ] || python3 -m virtualenv --quiet -p python3 ${virtualenv}
    source ${virtualenv}/bin/activate
    python3 -m pip install --quiet future protobuf numpy scikit-build scipy
    python3 -m pip install --quiet ninja pyyaml mkl mkl-include setuptools cmake cffi six typing
    python3 -m pip install --quiet --pre torch torchvision -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
    # python3 -m pip install --quiet ${third_party}/pytorch
    # python3 -m pip install --quiet ${third_party}/torchvision
    deactivate
}

schema() {
    bold "caffe2 schema"
    ${node_modules}/protobufjs/bin/pbjs -t static-module -w closure --no-encode --no-delimited --no-comments --no-convert --no-verify --no-create --keep-case --decode-text -r caffe2 -o ${src}/caffe2-proto.js ${third_party}/pytorch/caffe2/proto/caffe2.proto
    node ${tools}/update_pbjs.js enumeration ${src}/caffe2-proto.js floats float 1
}

metadata() {
    source ${virtualenv}/bin/activate
    pushd ${tools} > /dev/null
    bold "pytorch metadata"
    python3 pytorch-script.py metadata
    bold "caffe2 metadata"
    python3 caffe2-script.py metadata
    popd > /dev/null
    deactivate
}

zoo() {
    bold "pytorch zoo"
    source ${virtualenv}/bin/activate
    pushd ${tools} > /dev/null
    python3 pytorch-script.py zoo
    popd > /dev/null
    deactivate
}

while [ "$#" != 0 ]; do
    command="$1" && shift
    case "${command}" in
        "clean") clean;;
        "sync") sync;;
        "install") install;;
        "schema") schema;;
        "metadata") metadata;;
        "zoo") zoo;;
    esac
done
